package fun.coding.leetcode;

public class MedianOfTwoSortedArrays {

	public static void main(String[] args) {
		MedianOfTwoSortedArrays ins = new MedianOfTwoSortedArrays();
		int[] a1 = {1, 2};
		int[] b1 = {3, 4};
		System.out.println(ins.findMedianSortedArrays(a1, b1));
		
		int[] a2 = {1, 2, 3, 4};
		int[] b2 = {};
		System.out.println(ins.findMedianSortedArrays(a2, b2));
		
		int[] a3 = {2, 2, 2, 2};
		int[] b3 = {2, 2, 2};
		System.out.println(ins.findMedianSortedArrays(a3, b3));
		
		int[] a4 = {1, 2, 3, 4, 5};
		int[] b4 = {};
		System.out.println(ins.findMedianSortedArrays(a4, b4));
	}
	
	
	public double findMedianSortedArrays(int A[], int B[]) {
		int n = A.length + B.length;
		
		if (n % 2 == 0) {
			return (findKth(A, 0, B, 0, n / 2) + findKth(A, 0, B, 0, n / 2 + 1)) / 2.0;
		} else {
			return findKth(A, 0, B, 0, n / 2 + 1);
		}
		
	}
	
	public int findKth(int[] A, int A_start, int[] B, int B_start, int k){		
		if(A_start >= A.length) 
			return B[B_start + k - 1];
		if(B_start >= B.length)
			return A[A_start + k - 1];

		if (k == 1)
			return Math.min(A[A_start], B[B_start]);

		int A_key = A_start + k / 2 - 1 < A.length ? A[A_start + k / 2 - 1] : Integer.MAX_VALUE;
		int B_key = B_start + k / 2 - 1 < B.length ? B[B_start + k / 2 - 1] : Integer.MAX_VALUE; 

		if (A_key < B_key) {
			return findKth(A, A_start + k / 2, B, B_start, k - k / 2);
		} else {
			return findKth(A, A_start, B, B_start + k / 2, k - k / 2);
		}
	}
	
	
	// Ask clarification what is the median for even number?
	// If O(N) and using extra space, we could just keep 2 stacks, or one sorted list, directly get the median out, easy
	// The below implementation is O(N), but without any extra space
	public double findMedianSortedArraysLinear(int A[], int B[]) {
		if (A == null || B == null) return 0;
		
		int count = (A.length + B.length) / 2;
		boolean isEven = false;
		
		if ((A.length + B.length) % 2 == 0) {
			isEven = true;
		} else {
			count++;
		}
		
		int sum = 0;
		int i = 0;
		int j = 0; 
		int k = 1;
		
		while (i < A.length && j < B.length) {
			if (k == count && !isEven) {
				return Math.min(A[i], B[j]);
			} else if (k == count && isEven) {
				sum = Math.min(A[i], B[j]);
			} else if (k == count + 1 &&  isEven) {
				sum += Math.min(A[i], B[j]);
				return (double)sum/2;
			}
			
			if (A[i] < B[j]) {
				i++;
			} else {
				j++;
			}
			k++;
		}
		
		while (i < A.length) {
			if (k == count && !isEven) {
				return A[i];
			} else if (k == count && isEven) {
				sum = A[i];
			} else if (k == count + 1 && isEven) {
				sum += A[i];
				return (double)sum/2;
			}
			i++;
			k++;
		}
		
		while (j < B.length) {
			if (k == count && !isEven) {
				return B[j];
			} else if (k == count && isEven) {
				sum = B[j];
			} else if (k == count + 1 && isEven) {
				sum += B[j];
				return (double)sum/2;
			}
			j++;
			k++;
		}
		// not reach here
		return 0;
	}

}
